"""
@File  :get_npy_scp_from_dir.py
@Author:Xuelong Geng
@Date  :2024/6/15 16:17
@Desc  :
"""
import os

import tqdm
from gxl_ai_utils.utils import utils_file

input_dir = "/home/work_nfs14/xlgeng/asr_data_raw/wenetspeech4tts/TOKEN_4096-2/gxl_data"
npy_dict = utils_file.do_get_scp_for_wav_dir(input_dir, suffix='.npy')
output_dir = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list"
output_path = os.path.join(output_dir, "npy.scp")
utils_file.write_dict_to_scp(npy_dict, output_path)
wav_scp_path = "/home/work_nfs14/xlgeng/asr_data_raw/wenetspeech4tts/wav.scp"
text_scp_pat = "/home/work_nfs14/xlgeng/asr_data_raw/wenetspeech4tts/text"

wtn_list_path = os.path.join(output_dir, "wtn.list")

# get wtn.list
wav_dict = utils_file.load_dict_from_scp(wav_scp_path)
text_dict = utils_file.load_dict_from_scp(text_scp_pat)
wtn_list = []
for key in tqdm.tqdm(npy_dict.keys(), desc='get wtn list', total=len(npy_dict.keys())):
    temp_dict = {}
    key_wenet = key + '_wenet'
    temp_dict["key"] = key_wenet
    if key_wenet in wav_dict.keys():
        temp_dict["wav"] = wav_dict[key_wenet]
    else:
        print('not find', key_wenet, ' in wav.scp')
        continue
    if key_wenet in text_dict.keys():
        temp_dict["txt"] = text_dict[key_wenet]
    else:
        print('not find', key_wenet, ' in text.scp')
        continue
    temp_dict['npy'] = npy_dict[key]
    wtn_list.append(temp_dict)


utils_file.write_list_to_file(wtn_list, wtn_list_path)
